import math
import queue
import copy

import numpy as np
from config import *
import simpy

import uuid


class WCPNEnv:
    def __init__(self, agent):
        self.agent = agent
        self.N = USER_NODE_NUM  # 用户节点的数量
        self.M = MEC_NODE_NUM  # 边缘服务器的数量
        self.node_array = []
        self.task_in_channel_poll = set()  # 正在传输的任务列表
        self.sim_env = simpy.Environment()
        self.finish_work_set = set()        # 每一个时隙统计的完成的任务的
        self.reward_record = 0
        self.finish_task_num = 0

    # 4个边缘服务器
    def reset(self):
        print("[INFO] Reset Environment ...")
        # 重置
        self.sim_env = simpy.Environment()
        self.node_array = []
        self.task_in_channel_poll = set()  # 正在传输的任务列表
        self.sim_env = simpy.Environment()
        self.finish_work_set = set()        # 每一个时隙统计的完成的任务的
        self.reward_record = 0
        self.finish_task_num = 0
        # 先写固定的
        user_node1 = UserNode(0, 0.1*10**9, -150, 150, 0, 0.2)
        self.sim_env.process(user_node1.on_off_change(self.sim_env))
        self.node_array.append(user_node1)
        user_node2 = UserNode(1, 0.1*10**9, 150, 150, 0, 0.2)
        self.sim_env.process(user_node2.on_off_change(self.sim_env))
        self.node_array.append(user_node2)
        user_node3 = UserNode(2, 0.1*10**9, 450, 150, 0, 0.2)
        self.sim_env.process(user_node3.on_off_change(self.sim_env))
        self.node_array.append(user_node3)
        user_node4 = UserNode(3, 0.1*10**9, 150, 450, 0, 0.2)
        self.sim_env.process(user_node4.on_off_change(self.sim_env))
        self.node_array.append(user_node4)
        user_node5 = UserNode(4, 0.1*10**9, 150, -150, 0, 0.2)
        self.sim_env.process(user_node5.on_off_change(self.sim_env))
        self.node_array.append(user_node5)

        mec_node1 = MecNode(5, 50*10**9, 0, 0, 0, 2)
        self.node_array.append(mec_node1)
        mec_node2 = MecNode(6, 50*10**9, 300, 0, 0, 2)
        self.node_array.append(mec_node2)
        mec_node3 = MecNode(7, 50*10**9, 0, 300, 0, 2)
        self.node_array.append(mec_node3)
        mec_node4 = MecNode(8, 50*10**9, 300, 300, 0, 2)
        self.node_array.append(mec_node4)

        cloud_node = CloudNode()
        self.node_array.append(cloud_node)

    def test(self):
        task = Task(100000000, 5, self.sim_env.now)
        newTask = copy.deepcopy(task)
        mySet = set()
        mySet.add(task)
        mySet.add(newTask)
        mySet.remove(task)
        print(task)
        print(newTask)
        print(len(mySet))

    # 任务产生 采用on-off模型
    def generate(self, node_index):
        # print(f"[INFO] At Time: {self.sim_env.now} generating task...")
        node = self.node_array[node_index]
        if isinstance(node, UserNode):  # 如果是用户节点才可能产生任务
            if node.node_state == ON:   # 处于激活状态才产生任务
                is_generate = np.random.random()    # 随机数是否产生任务
                if is_generate < TASK_PROBABILITY:  # 产生任务
                    station = self.get_station(node.id)
                    # print(f"Node: {node.id} is generating task, station: {station}")
                    task = Task(100000000, 1, self.sim_env.now)
                    task.sour_id = node.id
                    # 计算任务本地计算的任务完成时间
                    node.task_to_send.add(task)
                    # print(f"[Task] {task} 任务生成并加入task_to_send")
                    node.task_generate = task
                    task.prev_state = self.get_current_state(node.id)
                    return True, task
        return False, None

    def cal_local_time(self, task):
        node = self.node_array[task.sour_id]
        sum_task = task.need_compute_time
        for task in node.task_queue:
            sum_task += task.need_compute_time
        return sum_task / node.compute_power

    # action 是一个一维矩阵 1*N 表示各个用户计算的任务卸载对象
    def step(self, action, node_index):
        # 调度
        self.transfer(node_index, action)
        self.node_array[node_index].task_generate = None

    # 调度 将节点i的任务调度到节点j
    # 调度到信道中
    def transfer(self, i, j):
        task_list = self.node_array[i].task_to_send  # 节点i的任务
        for task in task_list:
            task.action = j
            # print(f"[Task] {task} 调度")
            if task is not None:
                task.state = COMMUNICATING
                task.dest_id = j
                # 调度到MEC中
                if i != j:
                    if j >= USER_NODE_NUM:
                        # 不能直接和基站相连
                        if self.get_distance(i, j) >= STATION_COVER:
                            task.station_id = self.get_station(i)
                            # 基站j存在通信对象i
                            self.node_array[task.station_id].task_to_receive.add(task)
                            # self.node_array[task.station_id].user_set.add(i)
                        else:
                            task.station_id = j
                            self.node_array[j].task_to_receive.add(task)
                            # self.node_array[j].user_set.add(i)
                    else:
                        # 获得目的方的基站id
                        task.station_id = self.get_station(j)
                        if self.get_distance(i, task.station_id) >= STATION_COVER:
                            task.station_id = self.get_station(i)
                            # 基站j存在通信对象i
                            # self.node_array[task.station_id].user_set.add(i)
                            self.node_array[task.station_id].task_to_receive.add(task)
                        else:
                            #self.node_array[task.station_id].user_set.add(i)
                            self.node_array[task.station_id].task_to_receive.add(task)
                else:
                    task.station_id = i
                    self.node_array[i].task_to_receive.add(task)
                self.task_in_channel_poll.add(task)
                print(f"[Task] {task} transfer from {task.sour_id} to {task.dest_id} station is {task.station_id}")

    def communication(self):
        # print(f"[INFO] channel size: {len(self.task_in_channel_poll)}")
        for task in self.task_in_channel_poll:
            if task is not None:
                communicate_time = self.get_communicate_time(task)
                # 通信时间后到达节点 调用receive
                # print(task)
                # print("136 delay_receive")
                self.sim_env.process(self.delayed_receive(communicate_time, task, FIRST_TRANS))
        self.task_in_channel_poll = set()

    # 被调度的节点接收任务
    # 从信道到被卸载的节点
    ###################################################
    # 接受了之后怎么
    ###################################################
    def receive(self, task, flag):
        # if flag == FIRST_TRANS:
        #     self.task_in_channel_poll.remove(task)
        sour_id = task.sour_id
        station_id = task.station_id
        dest_id = task.dest_id
        # print(task)
        # print(f"send size: {len(self.node_array[sour_id].task_to_send)}" )
        # print(f"receive size: {len(self.node_array[station_id].task_to_receive)}")
        # print(sour_id)
        # print(station_id)
        # print(f"[Task] {task} 任务到达，从task_to_send中移除 {flag}")
        self.node_array[sour_id].task_to_send.remove(task)
        self.node_array[station_id].task_to_receive.remove(task)
        task.state = COMPUTING
        if self.node_array[dest_id].task_queue.qsize() == 0:
            self.node_array[dest_id].task_queue.put(task)
            compute_time = self.get_compute_time(task)
            self.sim_env.process(self.delayed_compute(compute_time, task))
        else:
            self.node_array[dest_id].task_queue.put(task)
        # print(self.node_array[dest_id].task_queue.qsize())
        # print(f"At Time: {self.sim_env.now} Node: {dest_id} task from {task.sour_id} is received, station is {task.station_id}")

    def compute(self, task):
        # 计算完成从队列中弹出
        self.node_array[task.dest_id].task_queue.get()
        # 计算下一个任务
        if self.node_array[task.dest_id].task_queue.qsize() != 0:
            next_task = self.node_array[task.dest_id].task_queue.queue[0]
            compute_time = self.get_compute_time(next_task)
            self.sim_env.process(self.delayed_compute(compute_time, next_task))

        task.finish_sub_num = task.finish_sub_num + 1
        # 周期性任务完成
        if task.finish_sub_num == task.sub_num_k:
            now = self.sim_env.now
            task.real_finish_time = now - task.generate_time
            self.finish_work_set.add(task)
            reward = -task.real_finish_time
            # print(f"[Task] {task} 任务完成, 时延:{now - task.generate_time}")
            # 任务完成 加入经验回放池
            self.agent.remember(task.prev_state, task.action, self.get_current_state(task.sour_id), reward)
            self.finish_task_num += 1
            self.reward_record += reward
            self.agent.train()

        # 未完成 发送下一个小任务
        else:
            # 进行发送子任务
            # new_sub_task = copy.deepcopy(task)
            new_sub_task = task
            self.node_array[task.sour_id].task_to_send.add(new_sub_task)
            self.node_array[task.station_id].task_to_receive.add(new_sub_task)
            communicate_time = self.get_communicate_time(new_sub_task)
            # 通信时间后到达节点 调用receive
            # print(new_sub_task)
            # print("185 delay_receive")
            self.sim_env.process(self.delayed_receive(communicate_time, new_sub_task, SECOND_TRANS))
            # print(f"[Task] {new_sub_task} 剩余任务数 {new_sub_task.sub_num_k-new_sub_task.finish_sub_num}单次任务完成，发送下一个子任务")

    # 根据每个station的任务个数分配带宽
    def allocate_bandwidth(self):
        for i in range(MEC_NODE_NUM):
            node = self.node_array[i+USER_NODE_NUM]
            node.bandwidth_for_user = node.bandwidth

    def get_current_state(self, node_index):
        agent_node = self.node_array[node_index]
        agent_state = np.zeros((1, AGENT_STATE_NUM))
        node_state = np.zeros((USER_NODE_NUM+MEC_NODE_NUM+1, NODE_STATE_NUM))
        if agent_node.task_generate is not None:
            k = agent_node.task_generate.sub_num_k
            t = agent_node.task_generate.need_compute_time
        else:
            k = 0
            t = 0
        communicate_power = agent_node.communicate_power
        compute_power = agent_node.compute_power
        wait_time = self.sum_queue_and_set(agent_node.task_queue, agent_node.task_to_send, agent_node.task_to_receive)
        wait_time = wait_time / agent_node.compute_power
        x1 = agent_node.x
        y1 = agent_node.y
        z1 = agent_node.z
        x2 = agent_node.x
        y2 = agent_node.y
        z2 = agent_node.z
        x3 = agent_node.x
        y3 = agent_node.y
        z3 = agent_node.z
        user_state = [k, t, communicate_power, compute_power, wait_time,
                      x1, y1, z1,
                      x2, y2, z2,
                      x3, y3, z3]
        agent_state[0] = user_state

        i = 0
        for node in self.node_array:
            if isinstance(node, UserNode):
                communication_power = node.compute_power
                compute_power = node.compute_power
                wait_time = self.sum_queue_and_set(node.task_queue, node.task_to_send, node.task_to_receive)
                wait_time = wait_time / node.compute_power
                x = node.x
                y = node.y
                z = node.z

                user_state = [communication_power, compute_power, wait_time,
                              x, y, z]
                node_state[i] = user_state
            else:
                bandwidth = node.bandwidth
                compute_power = node.compute_power
                x = node.x
                y = node.y
                z = node.z
                wait_time = self.sum_queue_and_set(node.task_queue, set(), node.task_to_receive)
                wait_time = wait_time / node.compute_power
                mec_state = [bandwidth, compute_power, wait_time,
                             x, y, z]
                node_state[i] = mec_state
            i += 1
        user_node_state_flat = agent_state.flatten()
        mec_node_state_flat = node_state.flatten()
        return np.concatenate((user_node_state_flat, mec_node_state_flat))

    def delayed_receive(self, delay, task, flag):
        yield self.sim_env.timeout(delay)
        self.receive(task, flag)
        # print(f"[INFO] Communication delay:{delay}")

    def delayed_compute(self, delay, task):
        yield self.sim_env.timeout(delay)
        self.compute(task)
        # print(f"[INFO] Compute delay:{delay}")

    def get_communicate_time(self, task):
        sour_id = task.sour_id
        dest_id = task.dest_id
        station_id = task.station_id
        if sour_id == dest_id:
            return 0
        if station_id == dest_id:
            return task.need_communicate_time / self.get_rij(sour_id, dest_id)
        else:
            return (task.need_communicate_time / self.get_rij(sour_id, station_id)
                    + task.need_communicate_time / self.get_rjj(station_id, dest_id))

    def get_compute_time(self, task):
        return task.need_compute_time / self.node_array[task.dest_id].compute_power

    def get_rij(self, i, j):
        B_j = self.node_array[j].bandwidth / len(self.node_array[j].task_to_receive)
        p_i = self.node_array[i].communicate_power / len(self.node_array[i].task_to_send)
        h_ij = 1
        alpha = 4   # 衰减指数，城市、郊区典型值 3-4
        N_0 = 10**(-20.4)  # 典型噪声功率
        I = self.get_I(i, j)
        r_ij = B_j*math.log((1+(p_i*h_ij*self.get_distance(i, j))/(N_0+I)), 2)
        return r_ij

    def get_rjj(self, j1, j2):
        # cable_length = 1000             # 电缆长度
        # max_distance = cable_length * 2
        # signal_speed = 200000000
        # attenuation = 1 - self.get_distance(j1, j2) / max_distance
        # return attenuation*signal_speed
        return 10**9 / self.get_distance(j1, j2)

    def get_I(self, i, j):
        I = 0
        for task in self.node_array[j].task_to_receive:
            if len(self.node_array[task.sour_id].task_to_send) != 0:
                I = I + self.node_array[task.sour_id].communicate_power / len(self.node_array[task.sour_id].task_to_send)
        return I

    def get_distance(self, i, j):
        node1 = self.node_array[i]
        node2 = self.node_array[j]
        distance = math.sqrt((node1.x - node2.x) ** 2 + (node1.y - node2.y) ** 2 + (node1.z - node2.z) ** 2)
        return distance

    # 传入用户节点的id, 返回节点连接的基站id
    def get_station(self, i):
        dist = self.get_distance(i, USER_NODE_NUM)
        station = USER_NODE_NUM
        for j in range(MEC_NODE_NUM):
            if dist > self.get_distance(i, j+USER_NODE_NUM):
                dist = self.get_distance(i, j+USER_NODE_NUM)
                station = j+USER_NODE_NUM
            elif dist == self.get_distance(i, j+USER_NODE_NUM):
                is_change = np.random.random()
                if is_change < 0.5:
                    dist = self.get_distance(i, j + USER_NODE_NUM)
                    station = j + USER_NODE_NUM
        return station

    def sum_queue_and_set(self, q, set1, set2):
        # 创建一个新队列，并深拷贝原队列中的元素
        return_sum = 0
        # 使用迭代器遍历原队列
        for item in list(q.queue):
            return_sum += item.need_compute_time
        for item in set1:
            return_sum += item.need_compute_time
        for item in set2:
            return_sum += item.need_compute_time
        return return_sum

    def calculate_reward(self):
        reward = 0
        for task in self.finish_work_set:
            reward -= task.real_finish_time
        self.finish_work_set = set()
        return reward


# 用户节点
class UserNode:
    def __init__(self, id, c, x, y, z, cp):
        self.id = id
        self.mec_id = -1                    # 用户设备连接到基站
        self.compute_power = c              # 算力大小
        self.x = x
        self.y = y
        self.z = z                          # 节点坐标
        self.communicate_power = cp         # 发射功率
        self.task_generate = None
        self.task_to_send = set()           # 产生的任务
        self.task_to_receive = set()  # 等待被接收的任务
        self.task_queue = queue.Queue()     # 待完成的任务
        self.node_state = OFF if np.random.random() < 0.5 else ON        # 初始状态
        # print(f"Node {self.id} state {self.node_state}")

    def on_off_change(self, env):
        while True:
            # print(f"Node {self.id} state {self.node_state}")
            delay = np.random.exponential(1 / lambda_on)
            # print(delay)
            self.node_state = OFF if self.node_state == ON else ON        # 初始状态
            yield env.timeout(delay)


# 边缘服务器节点
class MecNode:
    def __init__(self, id, c, x, y, z, b):
        self.id = id
        self.compute_power = c              # 算力大小
        self.x = x
        self.y = y
        self.z = z                          # 节点坐标
        self.bandwidth = b                  # 带宽
        self.bandwidth_for_user = b         # 分给每个用户的带宽
        self.task_to_receive = set()        # 等待被接收的任务
        self.task_queue = queue.Queue()     # 待完成的任务


# 边缘服务器节点
class CloudNode:
    def __init__(self):
        self.id = 0
        self.compute_power = 500*10**9      # 算力大小
        self.x = 10000
        self.y = 10000
        self.z = 10000 # 节点坐标
        self.bandwidth = 200  # 带宽
        self.task_to_receive = set()        # 等待被接收的任务
        self.task_queue = queue.Queue()     # 待完成的任务

class Task:
    def __init__(self, n, k, t):
        self.task_id = uuid.uuid4()
        self.sour_id = -1
        self.dest_id = -1           # 默认-1表示未被调度
        self.station_id = -1
        self.need_communicate_time = 1  # 需要通信任务的大小
        self.need_compute_time = n  # 需要CPU时间的大小
        self.sub_num_k = k          # 持续k次
        self.finish_sub_num = 0     # 已经完成的次数
        self.state = NEW            # 任务的状态 0 未被调度 1 被调度未到达 2 计算中 3 计算完成
        self.generate_time = t      # 任务生成的时间
        self.local_finish_time = 0
        self.real_finish_time = 0
        self.prev_state = None
        self.action = None
        # self.communicate_time = 0   # 通信时间


# # 某次产生的所有的任务
# class TaskAll:
#     def __init__(self):
#         self.N = USER_NODE_NUM
#         self.task_array = []
#
#     def generate(self):
#         for i in self.N:
#             # n = -1 表示没有
#             task = Task(10, 10)


